import sys
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QHBoxLayout, QLabel, QLineEdit, QTextEdit, QPushButton, \
    QComboBox, QFileDialog, QMenu, QAction
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QPalette, QBrush, QImage, QFont, QPixmap
from PyQt5.QtGui import QGuiApplication
from algorithms.caesar import caesar_encrypt, caesar_decrypt
from algorithms.base64 import base64_encode, base64_decode
from algorithms.aes import aes_encrypt, aes_decrypt, aes_encrypt_binary, aes_decrypt_binary
from algorithms.rsa import rsa_encrypt, rsa_decrypt, rsa_generate_keys, rsa_sign, rsa_verify
from algorithms.des import des_encrypt, des_decrypt
from algorithms.vigenere import vigenere_encrypt, vigenere_decrypt
from algorithms.railfence import railfence_encrypt, railfence_decrypt
from algorithms.ascii import ascii_encrypt, ascii_decrypt
from algorithms.morse import morse_encrypt, morse_decrypt
from algorithms.chacha import chacha_encrypt, chacha_decrypt
import os

class MainWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.rsa_private_key, self.rsa_public_key = rsa_generate_keys()

        # 设置背景图片
        self.set_background_image()

        # 设置Minecraft风格字体
        self.set_minecraft_font()

    def initUI(self):
        # 输入文本框
        self.input_label = QLabel('输入文本或选择文件内容:')
        self.input_text = QTextEdit()
        self.input_text.setContextMenuPolicy(3)  # 自定义上下文菜单
        self.input_text.customContextMenuRequested.connect(self.show_input_context_menu)

        # 密钥输入框 - 用于RSA签名/验证的公钥/私钥
        self.key_label = QLabel('输入密钥:')
        self.key_input = QLineEdit()

        # 添加密钥文件选择按钮
        self.select_key_file_button = QPushButton('选择密钥文件')
        self.select_key_file_button.clicked.connect(self.select_key_file)

        # 算法选择下拉框
        self.algorithm_label = QLabel('选择算法:')
        self.algorithm_combobox = QComboBox()
        self.algorithm_combobox.addItems(['凯撒密码', 'Base64', 'AES', 'RSA', 'DES', '维吉尼亚密码', '栅栏密码', 'ASCII', '摩斯密码','ChaCha'])

        # 操作选择下拉框 - 添加签名和验证选项
        self.operation_label = QLabel('选择操作:')
        self.operation_combobox = QComboBox()
        self.operation_combobox.addItems(['加密', '解密', '签名', '验证'])

        # 签名输入框 - 用于验证时输入签名
        self.signature_label = QLabel('输入签名:')
        self.signature_input = QLineEdit()
        self.signature_input.hide()  # 默认隐藏，只在验证时显示

        # 选择文件按钮
        self.select_file_button = QPushButton('选择明/密文文件')
        self.select_file_button.clicked.connect(self.select_file)

        # 生成RSA密钥对按钮
        self.generate_rsa_keys_button = QPushButton('生成RSA密钥对')
        self.generate_rsa_keys_button.clicked.connect(self.generate_rsa_keys)

        # 执行按钮 - Minecraft风格主按钮
        self.execute_button = QPushButton('执行')
        self.execute_button.setObjectName("mainButton")  # 设置对象名称以应用特殊样式
        self.execute_button.setStyleSheet("""
            padding: 5px;
            background-color: #4CAF50;
            color: #FFFFFF;
            border: 1px solid;
            border-color: #66BB6A #388E3C #388E3C #66BB6A;
            border-radius: 0;
            text-shadow: 2px 2px #000000;
            font-size: 16px;
        """)
        self.execute_button.clicked.connect(self.execute_operation)

        # 输出文本框
        self.output_label = QLabel('输出结果:')
        self.output_text = QTextEdit()
        self.output_text.setReadOnly(True)
        self.output_text.setContextMenuPolicy(3)  # 自定义上下文菜单
        self.output_text.customContextMenuRequested.connect(self.show_output_context_menu)

        # 布局
        input_layout = QVBoxLayout()
        input_layout.addWidget(self.input_label)
        input_layout.addWidget(self.input_text)

        # 更新密钥布局，添加密钥文件选择按钮
        key_layout = QHBoxLayout()
        key_layout.addWidget(self.key_label)
        key_layout.addWidget(self.key_input)
        key_layout.addWidget(self.select_key_file_button)

        algorithm_layout = QHBoxLayout()
        algorithm_layout.addWidget(self.algorithm_label)
        algorithm_layout.addWidget(self.algorithm_combobox)

        operation_layout = QHBoxLayout()
        operation_layout.addWidget(self.operation_label)
        operation_layout.addWidget(self.operation_combobox)

        # 添加签名输入框布局
        signature_layout = QHBoxLayout()
        signature_layout.addWidget(self.signature_label)
        signature_layout.addWidget(self.signature_input)

        # 修改按钮布局，将生成RSA密钥对按钮放在执行按钮左边
        button_layout = QHBoxLayout()
        button_layout.addWidget(self.select_file_button)
        button_layout.addWidget(self.generate_rsa_keys_button)
        button_layout.addWidget(self.execute_button)

        output_layout = QVBoxLayout()
        output_layout.addWidget(self.output_label)
        output_layout.addWidget(self.output_text)

        main_layout = QVBoxLayout()
        main_layout.addLayout(input_layout)
        main_layout.addLayout(key_layout)
        main_layout.addLayout(algorithm_layout)
        main_layout.addLayout(operation_layout)
        main_layout.addLayout(signature_layout)  # 添加签名输入框布局
        main_layout.addLayout(button_layout)
        main_layout.addLayout(output_layout)

        self.setLayout(main_layout)
        self.setWindowTitle('简易密码机')
        self.setGeometry(300, 300, 800, 600)

        # 连接操作选择变化事件
        self.operation_combobox.currentTextChanged.connect(self.update_ui_for_operation)
        self.update_ui_for_operation(self.operation_combobox.currentText())

    def set_background_image(self):
        """设置背景图片"""
        # 请将 'path_to_your_image.jpg' 替换为你实际的图片路径
        # pixmap = QPixmap('D:/1astudyWork/密码实训/crypto-system/crypto-tool/crypto-tool/src/ui/background.png')
        # pixmap = QPixmap('F:/crypto-tool/crypto-tool/src/ui/background.png')
        # 获取当前脚本所在的目录
        script_dir = os.path.dirname(os.path.abspath(__file__))
        # 构建背景图片的相对路径
        image_path = os.path.join(script_dir, 'background.png')
        pixmap = QPixmap(image_path)
        if not pixmap.isNull():
            # 获取窗口的大小
            window_size = self.size()
            # 缩放图片以适应窗口大小
            scaled_pixmap = pixmap.scaled(window_size, Qt.KeepAspectRatioByExpanding, Qt.SmoothTransformation)
            palette = self.palette()
            palette.setBrush(QPalette.Window, QBrush(scaled_pixmap))
            self.setPalette(palette)
            self.setAutoFillBackground(True)
        else:
            print("图片加载失败，请检查main_window.py中set_background_image函数下的pixmap路径")

    def set_minecraft_font(self):
        """设置Minecraft风格字体"""
        # 尝试加载Minecraft风格字体，如果没有则使用默认的等宽字体
        font = QFont("Minecraft", 12)
        font.setStyleStrategy(QFont.PreferAntialias)
        self.setFont(font)

    def update_ui_for_operation(self, operation):
        """根据选择的操作更新UI显示"""
        algorithm = self.algorithm_combobox.currentText()

        # 仅当选择RSA算法时启用签名/验证选项
        if algorithm == 'RSA' and operation in ['签名', '验证']:
            self.key_input.setEnabled(True)
            self.select_key_file_button.setEnabled(True)

            if operation == '签名':
                self.signature_input.hide()
                self.signature_label.hide()
                self.key_label.setText('输入私钥:')
                self.key_input.setText(self.rsa_private_key.decode('utf-8'))
            else:  # 验证
                self.signature_input.show()
                self.signature_label.show()
                self.key_label.setText('输入公钥:')
                self.key_input.setText(self.rsa_public_key.decode('utf-8'))
        else:
            self.signature_input.hide()
            self.signature_label.hide()
            self.key_label.setText('输入密钥:')
            self.select_key_file_button.setEnabled(True)

            # 恢复其他算法的默认密钥处理
            if algorithm == 'RSA' and operation in ['加密', '解密']:
                if operation == '加密':
                    self.key_input.setText(self.rsa_public_key.decode('utf-8'))
                else:
                    self.key_input.setText(self.rsa_private_key.decode('utf-8'))
            # 对于不需要密钥的算法，禁用密钥输入
            elif algorithm == 'Base64':
                self.key_input.setEnabled(False)
                self.select_key_file_button.setEnabled(False)
                self.key_input.setText('')

    def select_file(self):
        file_path, _ = QFileDialog.getOpenFileName(self, '选择文件', '', '文本文件 (*.txt);;图片文件 (*.png *.jpg *.jpeg);;视频文件 (*.mp4 *.avi);;加密文件 (*.enc)')
        if file_path:
            try:
                if file_path.lower().endswith(('.png', '.jpg', '.jpeg', '.mp4', '.avi', '.enc')):
                    with open(file_path, 'rb') as file:
                        self.input_binary = file.read()
                        self.input_text.setPlainText(f"已选择二进制文件: {os.path.basename(file_path)}")
                else:
                    with open(file_path, 'r', encoding='utf-8') as file:
                        content = file.read()
                        self.input_text.setPlainText(content)
            except Exception as e:
                print(f"读取文件出错: {e}")

    def select_key_file(self):
        # 修改文件类型过滤器，添加 .pem 文件类型
        file_path, _ = QFileDialog.getOpenFileName(self, '选择密钥文件', '', '文本文件 (*.txt *.pem);;所有文件 (*)')
        if file_path:
            try:
                with open(file_path, 'r', encoding='utf-8') as file:
                    key_content = file.read()
                    self.key_input.setText(key_content)
            except Exception as e:
                error_msg = f"读取密钥文件出错: {e}"
                self.output_text.setPlainText(error_msg)
                print(error_msg)

    def execute_operation(self):
        input_text = self.input_text.toPlainText()
        key = self.key_input.text()
        algorithm = self.algorithm_combobox.currentText()
        operation = self.operation_combobox.currentText()
        signature = self.signature_input.text()

        if hasattr(self, 'input_binary'):
            if algorithm == 'AES':
                if len(key) != 16 and len(key) != 24 and len(key) != 32:
                    result = "AES密钥长度必须为16、24或32字节"
                else:
                    if operation == '加密':
                        encrypted_data = aes_encrypt_binary(self.input_binary, key)
                        save_path, _ = QFileDialog.getSaveFileName(self, '保存加密文件', '', '加密文件 (*.enc)')
                        if save_path:
                            with open(save_path, 'wb') as file:
                                file.write(encrypted_data)
                            result = f"加密文件已保存到 {save_path}"
                        else:
                            result = "取消保存"
                    elif operation == '解密':
                        decrypted_data = aes_decrypt_binary(self.input_binary, key)
                        save_path, _ = QFileDialog.getSaveFileName(self, '保存解密文件', '', '图片文件 (*.png *.jpg *.jpeg);;视频文件 (*.mp4 *.avi)')
                        if save_path:
                            with open(save_path, 'wb') as file:
                                file.write(decrypted_data)
                            result = f"解密文件已保存到 {save_path}"
                        else:
                            result = "取消保存"
            else:
                result = "仅支持使用AES算法对图片和视频进行加解密"
        else:
            if algorithm == '凯撒密码':
                try:
                    key = int(key)
                    if operation == '加密':
                        result = caesar_encrypt(input_text, key)
                    else:
                        result = caesar_decrypt(input_text, key)
                except ValueError:
                    result = "密钥必须是整数"
            elif algorithm == 'Base64':
                if operation == '加密':
                    result = base64_encode(input_text)
                else:
                    result = base64_decode(input_text)
            elif algorithm == 'AES':
                if len(key) != 16 and len(key) != 24 and len(key) != 32:
                    result = "AES密钥长度必须为16、24或32字节"
                else:
                    if operation == '加密':
                        result = aes_encrypt(input_text, key)
                    else:
                        result = aes_decrypt(input_text, key)
            elif algorithm == 'RSA':
                if operation == '加密':
                    result = rsa_encrypt(input_text, key)
                elif operation == '解密':
                    result = rsa_decrypt(input_text, key)
                elif operation == '签名':
                    result = rsa_sign(input_text, key)
                elif operation == '验证':
                    result = rsa_verify(input_text, signature, key)
            elif algorithm == 'DES':
                if len(key) != 8:
                    result = "DES密钥长度必须为8字节"
                else:
                    if operation == '加密':
                        result = des_encrypt(input_text, key)
                    else:
                        result = des_decrypt(input_text, key)
            elif algorithm == '维吉尼亚密码':
                if not key.isalpha():
                    result = "维吉尼亚密码密钥必须由字母组成"
                else:
                    if operation == '加密':
                        result = vigenere_encrypt(input_text, key)
                    else:
                        result = vigenere_decrypt(input_text, key)
            elif algorithm == '栅栏密码':
                try:
                    key_int = int(key)
                    if key_int <= 1:
                        result = "栅栏密码密钥必须大于1"
                    else:
                        if operation == '加密':
                            result = railfence_encrypt(input_text, key_int)
                        else:
                            result = railfence_decrypt(input_text, key_int)
                except ValueError:
                    result = "栅栏密码密钥必须是整数"
            elif algorithm == 'ASCII':
                if operation == '加密':
                    result = ascii_encrypt(input_text)
                else:
                    result = ascii_decrypt(input_text)
            elif algorithm == '摩斯密码':
                if operation == '加密':
                    result = morse_encrypt(input_text)
                else:
                    result = morse_decrypt(input_text)
            elif algorithm == 'ChaCha':
                if len(key) != 32:
                    result = "ChaCha密钥长度必须为32字节"
                else:
                    if operation == '加密':
                        result = chacha_encrypt(input_text, key)
                    else:
                        result = chacha_decrypt(input_text, key)

        self.output_text.setPlainText(result)

    # 添加缺失的上下文菜单处理方法
    def show_input_context_menu(self, pos):
        menu = QMenu(self)
        paste_action = QAction('粘贴', self)
        paste_action.triggered.connect(self.paste_to_input)
        menu.addAction(paste_action)
        menu.exec_(self.input_text.mapToGlobal(pos))

    def show_output_context_menu(self, pos):
        menu = QMenu(self)
        copy_action = QAction('复制', self)
        copy_action.triggered.connect(self.copy_from_output)
        menu.addAction(copy_action)
        menu.exec_(self.output_text.mapToGlobal(pos))

    def paste_to_input(self):
        clipboard = QApplication.clipboard()
        text = clipboard.text()
        self.input_text.setPlainText(text)

    def copy_from_output(self):
        clipboard = QApplication.clipboard()
        text = self.output_text.toPlainText()
        clipboard.setText(text)

    def generate_rsa_keys(self):
        private_key, public_key = rsa_generate_keys()

        private_key_path, _ = QFileDialog.getSaveFileName(self, '保存私钥文件', '', 'PEM文件 (*.pem)')
        if private_key_path:
            with open(private_key_path, 'wb') as file:
                file.write(private_key)

        public_key_path, _ = QFileDialog.getSaveFileName(self, '保存公钥文件', '', 'PEM文件 (*.pem)')
        if public_key_path:
            with open(public_key_path, 'wb') as file:
                file.write(public_key)

        result = f"私钥已保存到 {private_key_path}\n公钥已保存到 {public_key_path}"
        self.output_text.setPlainText(result)
